//============================
// Stereo Vision Sandbox
// Feb 2014 by Isop Alexander
// [isop.alexander@gmail.com]
//============================

#include "cv.h"
#include "highgui.h"


// Define function headers
void combineTwoImages(cv::Mat &dst, const cv::Mat img1, const cv::Mat img2);
IplImage* barrel_pincusion_dist(IplImage* img, double kx, double ky);
cv::Mat stereo_vision(IplImage* img_left, IplImage* img_right, CvSize target_size, double Kx, double Ky);


int main(int argc, char* argv[])
{
	// Init windows to show image data
	cv::namedWindow("Display Stereo", cv::WINDOW_AUTOSIZE);


	// Capture image from camera
	//---------------------------
	//CvCapture* capture = cvCaptureFromCAM(0);
	//IplImage* img = cvQueryFrame(capture);

	//while(1)
	//{
	//	img=cvQueryFrame(capture);

	//	cv::Mat mat_stereo=stereo_vision(img, img, cvSize(1280, 480), 0.000001, 0.000001); // Combine to stereo vision image
	//	cv::imshow("Display Stereo", mat_stereo); // Display iamge

	//	cvWaitKey(50);
	//}


	// Use one single image to show stereo vision
	//--------------------------------------------
	IplImage* img = cvLoadImage("C:\\image.jpg"); // Load as IplImage
	cv::Mat mat_stereo = stereo_vision(img, img, cvSize(1280, 480), 0.000001, 0.000001); // Combine to stereo vision image
	cv::imshow("Display Stereo", mat_stereo); // Display iamge
	cv::waitKey(0); // Wait for user to press any key

	// Clean up
	cvReleaseImage(&img);

	return 0;
}


// Combine two distorted images to one image (e.g. 2x 640x460 into 1x 1280x480)
cv::Mat stereo_vision(IplImage* img_left, IplImage* img_right, CvSize target_size, double Kx, double Ky)
{
	cv::Mat mat_final(target_size.width, target_size.height, CV_8UC3, cv::Scalar(0, 0, 255));
	cv::Mat mat_left_distorted(img_left->width, img_left->width, CV_8UC3, cv::Scalar(0, 0, 255));
	cv::Mat mat_right_distorted(img_right->width, img_right->width, CV_8UC3, cv::Scalar(0, 0, 255));

	IplImage* img_left_distorted = cvCloneImage(img_left);
	IplImage* img_right_distorted = cvCloneImage(img_right);

	img_left_distorted = barrel_pincusion_dist(img_left, Kx, Ky);
	img_right_distorted = barrel_pincusion_dist(img_right, Kx, Ky);

	mat_left_distorted = cv::cvarrToMat(img_left_distorted, true);
	mat_right_distorted = cv::cvarrToMat(img_right_distorted, true);

	combineTwoImages(mat_final, mat_left_distorted, mat_right_distorted);

	return mat_final;
}

void combineTwoImages(cv::Mat &dst, const cv::Mat img1, const cv::Mat img2)
{
	if (img1.cols != img2.cols || img1.rows != img2.rows) {
		std::cout << "img1 and img2 should be identical size";
	}
	int rows = img1.rows;
	int cols = img1.cols;
	dst = cvCreateMat(rows, 2 * cols, img1.type());
	cv::Mat tmp = dst(cv::Rect(0, 0, cols, rows));
	img1.copyTo(tmp);
	tmp = dst(cv::Rect(cols, 0, cols, rows));
	img2.copyTo(tmp);
}


// Main function to apply pincussion distortion, also refer to:
// http://stackoverflow.com/questions/6199636/formulas-for-barrel-pincushion-distortion
IplImage* barrel_pincusion_dist(IplImage* img, double kx, double ky)
{
	IplImage* mapx = cvCreateImage(cvGetSize(img), IPL_DEPTH_32F, 1);
	IplImage* mapy = cvCreateImage(cvGetSize(img), IPL_DEPTH_32F, 1);

	int w = img->width;
	int h = img->height;

	double Cx = double(w / 2);
	double Cy = double(h / 2);

	float* pbuf = (float*)mapx->imageData;
	for (int y = 0; y < h; y++)
	{
		for (int x = 0; x < w; x++)
		{
			float u = Cx + (x - Cx)*(1 + kx*((x - Cx)*(x - Cx) + (y - Cy)*(y - Cy)));
			*pbuf = u;
			++pbuf;
		}
	}

	pbuf = (float*)mapy->imageData;
	for (int y = 0; y < h; y++)
	{
		for (int x = 0; x < w; x++)
		{
			*pbuf = Cy + (y - Cy)*(1 + ky*((x - Cx)*(x - Cx) + (y - Cy)*(y - Cy)));
			++pbuf;
		}
	}

	IplImage* temp = cvCloneImage(img);
	cvRemap(temp, img, mapx, mapy);
	cvReleaseImage(&temp);
	cvReleaseImage(&mapx);
	cvReleaseImage(&mapy);

	return img;
}